import os
import glob
import random
import argparse
from PIL import Image
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt


def find_image(uuid, image_dir):
    # Find the image file for the given UUID (non-sorted, ident-only, BMP)
    pattern = os.path.join(image_dir, f"{uuid}~*~origin~ident~512.bmp")
    files = glob.glob(pattern)
    if files:
        return files[0]
    return None


def images_to_pdf(image_info_list, pdf_path, images_per_row=4, images_per_page=32):
    # image_info_list: list of (image_path, cluster_label, uuid, family_name)
    with PdfPages(pdf_path) as pdf:
        for i in range(0, len(image_info_list), images_per_page):
            fig, axes = plt.subplots(
                nrows=images_per_page // images_per_row,
                ncols=images_per_row,
                figsize=(images_per_row * 3, (images_per_page // images_per_row) * 3)
            )
            axes = axes.flatten()
            for j, ax in enumerate(axes):
                idx = i + j
                ax.axis('off')
                if idx < len(image_info_list):
                    image_path, cluster_label, uuid, family_name = image_info_list[idx]
                    try:
                        img = Image.open(image_path)
                        ax.imshow(img)
                        # Show UUID, family name, and cluster label
                        ax.set_title(f"{uuid} ({family_name})\nCluster {cluster_label}", fontsize=8)
                    except Exception as e:
                        ax.set_title(f"Error: {e}", fontsize=8)
            plt.tight_layout()
            pdf.savefig(fig)
            plt.close(fig)


def process_clusters(clustered_dir, image_dir, max_per_cluster=100):
    for fname in os.listdir(clustered_dir):
        if fname.startswith('cluster_') and fname.endswith('.txt'):
            cluster_path = os.path.join(clustered_dir, fname)
            # Extract cluster label from filename
            cluster_label = int(fname[len('cluster_'):-len('.txt')])
            with open(cluster_path) as f:
                uuids = [line.strip() for line in f if line.strip()]
            if not uuids:
                continue
            selected = random.sample(uuids, min(max_per_cluster, len(uuids)))
            image_info_list = []
            for uuid in selected:
                image_path = find_image(uuid, image_dir)
                if image_path:
                    # Extract family name from BMP file name (second field delimited by '~')
                    base = os.path.basename(image_path)
                    parts = base.split('~')
                    family_name = parts[1] if len(parts) > 1 else ''
                    image_info_list.append((image_path, cluster_label, uuid, family_name))
            if not image_info_list:
                continue
            pdf_path = os.path.join(clustered_dir, fname.replace('.txt', '_report.pdf'))
            print(f"Generating PDF for {fname} with {len(image_info_list)} images...")
            images_to_pdf(image_info_list, pdf_path)


def process_non_dominant(non_dom_dir, image_dir, max_per_cluster=20):
    for fname in os.listdir(non_dom_dir):
        if fname.startswith('family_') and fname.endswith('.md'):
            fam_path = os.path.join(non_dom_dir, fname)
            with open(fam_path) as f:
                lines = [line.rstrip() for line in f]
            cluster_blocks = {}
            current_cluster = None
            for line in lines:
                if line.startswith('## Cluster '):
                    current_cluster = int(line.split('## Cluster ')[1].replace(':', ''))
                    cluster_blocks[current_cluster] = []
                elif line.strip() and current_cluster is not None:
                    cluster_blocks[current_cluster].append(line.strip())
            all_image_info = []
            for cl, uuids in cluster_blocks.items():
                selected = random.sample(uuids, min(max_per_cluster, len(uuids)))
                for uuid in selected:
                    image_path = find_image(uuid, image_dir)
                    if image_path:
                        all_image_info.append((image_path, cl, uuid))
            if not all_image_info:
                continue
            pdf_path = os.path.join(non_dom_dir, fname.replace('.md', '_report.pdf'))
            print(f"Generating PDF for {fname} with {len(all_image_info)} images...")
            images_to_pdf(all_image_info, pdf_path)


def main():
    parser = argparse.ArgumentParser(description='Generate PDF reports for clusters and non-dominant families.')
    parser.add_argument('--analysis_dir', type=str, required=True, help='Directory containing analysis results')
    parser.add_argument('--image_dir', type=str, required=True, help='Directory containing image files')
    parser.add_argument('--seed', type=int, default=42, help='Random seed')
    args = parser.parse_args()
    random.seed(args.seed)

    clustered_dir = os.path.join(args.analysis_dir, 'clustered_instances')
    non_dom_dir = os.path.join(args.analysis_dir, 'non_dominant_families')

    print('Processing clusters...')
    process_clusters(clustered_dir, args.image_dir)
    print('Processing non-dominant families...')
    # process_non_dominant(non_dom_dir, args.image_dir)
    print('Done.')

if __name__ == '__main__':
    main() 